#include "stdafx.h"
#include "ogglength.h"
#include "lhcutils/IStream.h"
#include "lhcutils/Exceptions.h"
#include "OggVorbisException.h"
#include "lhcutils/utilities.h"
#include "lhcutils/common_macros.h"
#include <vorbis/vorbisfile.h>
#include <vorbis/codec.h>
#include <boost/shared_ptr.hpp>
#include <boost/optional.hpp>
#include <boost/scope_exit.hpp>
#include <vector>
#include <cerrno>

using namespace std;
using namespace boost;
using namespace lhcutils;

namespace ogglength
{

StringPtr vorbisCorruptMsg = ustr("The file does not appear to be an Ogg Vorbis file.");

ov_callbacks GetOvStreamCallbacks()
{
	ov_callbacks callbacks;
	callbacks.close_func = NULL;
	callbacks.read_func = StreamReadFunc;
	callbacks.seek_func = StreamSeekFunc;
	callbacks.tell_func = StreamTellFunc;

	return callbacks;
}

struct OvCallbackStream
{
public:
	IStream* OggStream;
	boost::optional<IoException> LastIoError;
	boost::optional<Exception> LastPlainError;
	OvCallbackStream(IStream* streamPtr)
	{
		OggStream = streamPtr;
	}
};

size_t StreamReadFunc(void* buffer, size_t size, size_t count, void* dataSource)
{
	OvCallbackStream* streamPtr;

	// libvorbisfile expects 0 errno on return of 0 meaning EOF,
	// non-0 errno on return of 0 meaning error
	errno = 0;

	try
	{
		streamPtr = reinterpret_cast<OvCallbackStream*>(dataSource);
		size_t bytesRead = streamPtr->OggStream->Read(buffer, 0, size * count); // Possible overflow when multiplying, but shouldn't matter
		return bytesRead;
	}
	catch(const IoException& ex)
	{
		errno = 1;
		streamPtr->LastIoError = ex;
	}
	catch(const Exception& ex)
	{
		errno = 1;
		streamPtr->LastPlainError = ex;
		return 0;
	}
	catch(...)
	{
		errno = 1;
		return 0;
	}
}

int StreamSeekFunc(void* dataSource, ogg_int64_t offset, int whence)
{
	OvCallbackStream* streamPtr;
	errno = 0;

	try
	{
		streamPtr = reinterpret_cast<OvCallbackStream*>(dataSource);
		SeekOrigin::Value origin;
		if(whence == SEEK_SET)
			origin = SeekOrigin::Begin;
		else if(whence == SEEK_CUR)
			origin = SeekOrigin::Current;
		else
			origin = SeekOrigin::End;

		streamPtr->OggStream->SetPosition(offset, origin);
		return 0;
	}
	catch(const IoException& ex)
	{
		errno = 1;
		streamPtr->LastIoError = ex;
	}
	catch(const Exception& ex)
	{
		errno = 1;
		streamPtr->LastPlainError = ex;
		return -1;
	}
	catch(...)
	{
		errno = 1;
		return -1;
	}
}

long StreamTellFunc(void* dataSource)
{
	OvCallbackStream* streamPtr;
	errno = 0;

	try
	{
		streamPtr = reinterpret_cast<OvCallbackStream*>(dataSource);
		int64_t position = streamPtr->OggStream->GetPosition();
		return position;
	}
	catch(const IoException& ex)
	{
		errno = 1;
		streamPtr->LastIoError = ex;
	}
	catch(const Exception& ex)
	{
		errno = 1;
		streamPtr->LastPlainError = ex;
		return -1;
	}
	catch(...)
	{
		errno = 1;
		return -1;
	}
}

void ThrowOvOpenException(OvCallbackStream& streamCallbackData, int openResult);

double GetReportedTime(IStream& oggStream)
{
	if(!oggStream.CanRead())
	{
		throw ProgrammerErrorException(ustr("Stream is not readable."));
	}
	if(!oggStream.CanSeek())
	{
		throw ProgrammerErrorException(ustr("Streamis not seekable."));
	}
	
	ov_callbacks callbacks = GetOvStreamCallbacks();
	OggVorbis_File vfStruct;
	OvCallbackStream streamCallbackData(&oggStream);

	int openResult = ov_open_callbacks(&streamCallbackData, &vfStruct, NULL, NULL, callbacks);

	if(openResult != 0)
	{
		ThrowOvOpenException(streamCallbackData, openResult);
	}

	BOOST_SCOPE_EXIT( (&vfStruct) )
	{
		UNUSED int ret = ov_clear(&vfStruct);
	} BOOST_SCOPE_EXIT_END

	// -1 = get time for entire physical bitstream, not just one logical bitstream.
	// Files we care about should only have 1 logical bitstream anyway.
	double reportedTime = ov_time_total(&vfStruct, -1);

	if(reportedTime == OV_EINVAL)
	{
		throw OggVorbisException(ustr("Error getting reported length."));
	}

	return reportedTime;
}

void ThrowCallbackExceptionIfPresent(OvCallbackStream& streamCallbackData)
{
	if(streamCallbackData.LastIoError.is_initialized())
	{
		throw streamCallbackData.LastIoError.get();
	}
	else if(streamCallbackData.LastPlainError.is_initialized())
	{
		throw streamCallbackData.LastPlainError.get();
	}
}

void ThrowOvOpenException(OvCallbackStream& streamCallbackData, int openResult)
{
	ThrowCallbackExceptionIfPresent(streamCallbackData);

	StringPtr vorbisErrorText;
	switch(openResult)
	{
	case OV_EREAD:
		throw IoException(str("Error reading.")); // Can this ever be reached?
		break;
	case OV_ENOTVORBIS:
		vorbisErrorText = vorbisCorruptMsg;
		break;
	case OV_EVERSION:
		vorbisErrorText = ustr("The file uses an unsupported version of Ogg Vorbis.");
		break;
	case OV_EBADHEADER:
		vorbisErrorText = vorbisCorruptMsg;
		break;
	case OV_EFAULT:
		vorbisErrorText = ustr("There was an internal problem.");
		break;
	default:
		vorbisErrorText = ustr("Unknown error.");
	}

	throw OggVorbisException(vorbisErrorText);
}

void ThrowOvReadException(OvCallbackStream& streamCallbackData, long readReturn)
{
	ThrowCallbackExceptionIfPresent(streamCallbackData);

	StringPtr vorbisErrorText;
	switch(readReturn)
	{
	case OV_HOLE:
		// indicates there was an interruption in the data.
		//(one of: garbage between pages, loss of sync followed by recapture, or a corrupt page)
	case OV_EBADLINK:
		// indicates that an invalid stream section was supplied to libvorbisfile, or the requested link is corrupt.
	case OV_EINVAL:
		// indicates the initial file headers couldn't be read or are corrupt, or that the initial open call for vf failed.
		vorbisErrorText = vorbisCorruptMsg;
		break;
	default:
		vorbisErrorText = ustr("Unknown error.");
	}

	throw OggVorbisException(vorbisErrorText);
}

double GetRealTime(IStream& oggStream)
{
	// Get the real song length by decoding the vorbis stream and doing the math with the number of samples
	// each ov_read gives.

	if(!oggStream.CanRead())
	{
		throw ProgrammerErrorException(ustr("Stream is not readable."));
	}
	if(!oggStream.CanSeek())
	{
		throw ProgrammerErrorException(ustr("Stream is not seekable."));
	}

	ov_callbacks callbacks = GetOvStreamCallbacks();
	OggVorbis_File vfStruct;
	OvCallbackStream streamCallbackData(&oggStream);

	int openResult = ov_open_callbacks(&streamCallbackData, &vfStruct, NULL, NULL, callbacks);

	if(openResult != 0)
	{
		ThrowOvOpenException(streamCallbackData, openResult);
	}

	BOOST_SCOPE_EXIT( (&vfStruct) )
	{
		UNUSED int ret = ov_clear(&vfStruct);
	} BOOST_SCOPE_EXIT_END
	
	double totalTimeRead = 0; // in seconds
	char buffer[4096];
	int logicalBitstreamRead = -555; // Number of the logical bitstream of the current page
	
	// Number of the logical bitstream we've been reading the whole time, -1 if we're on the first page.
	// If different from logicalBistreamRead, there's more than one logical bitstream so we bail.
	int logicalBitstreamReading = -1;

	long bytesRead;
	int pcmWordSize = 2; // 16-bit samples
	while((bytesRead = ov_read(&vfStruct, buffer, 4096, 0, pcmWordSize, 1, &logicalBitstreamRead)) > 0)
	{
		if(logicalBitstreamReading == -1)
		{
			// First Ogg page - this should be the logical bitstream all pages are in.
			logicalBitstreamReading = logicalBitstreamRead;
		}
		else if(logicalBitstreamReading != logicalBitstreamRead)
		{
			// A page in a logical bitstream different from the one we've been reading
			throw OggVorbisException(ustr("There is more than one logical bitstream in the file. I don't know how to handle that."));
		}
		long samplesRead = bytesRead / pcmWordSize; // actually this is samples * channels
		int numChannels = vfStruct.vi[logicalBitstreamRead].channels;
		if(numChannels <= 0) // Don't crash with a divide by 0
		{
			throw OggVorbisException(vorbisCorruptMsg);
		}
		long samplesPerChannel = samplesRead / numChannels;
		double timeRead = static_cast<double>(samplesPerChannel) / vfStruct.vi[logicalBitstreamRead].rate;

		totalTimeRead += timeRead;
	}

	if(bytesRead < 0)
	{
		ThrowOvReadException(streamCallbackData, bytesRead);
	}

	return totalTimeRead;
}

void ChangeSongLength(IStream& oggStream, double numSeconds)
{
	// For details of the Ogg format, see http://xiph.org/ogg/doc/, http://xiph.org/ogg/doc/oggstream.html,
	// http://xiph.org/ogg/doc/framing.html, http://en.wikipedia.org/wiki/Ogg
	//
	// For details of the Vorbis format, see http://xiph.org/vorbis/doc/Vorbis_I_spec.html
	
	if(!oggStream.CanRead())
	{
		throw ProgrammerErrorException(ustr("Stream is not readable."));
	}
	if(!oggStream.CanSeek())
	{
		throw ProgrammerErrorException(ustr("Stream is not seekable."));
	}
	if(!oggStream.CanWrite())
	{
		throw ProgrammerErrorException(ustr("Stream is not writeable."));
	}

	uint8_t version; // Version field of an Ogg page
	uint8_t headerType; // Header type field of an Ogg page.
	ogg_uint32_t sampleRate = 0;
	ogg_int32_t savedBitstreamSerialNumber = 0; // assignment stops a gcc warning. I know that it will have been initialized when it's used.

	// Read Ogg pages until we get to the last page (indicated by the "end of stream" bit set in the Ogg page header).
	// Have the seek pointer at granule position then.
	// Special care if given to the first page, because that is the primary Vorbis header and contains
	// the sample rate, which is needed to calculate what we should set the granule position of the last page to.
	while(true)
	{
		// All Ogg pages begin with the bytes "OggS"
		vector<uint8_t> buffer = oggStream.ReadFull(4);
		if(buffer[0] != 'O' || buffer[1] != 'g' || buffer[2] != 'g' || buffer[3] != 'S')
		{
			throw OggVorbisException(vorbisCorruptMsg);
		}

		// Ogg version field. Currently should always be 0.
		version = oggStream.ReadByte();
		if(version != 0)
		{
			throw OggVorbisException(ustr("The file uses an unsupported version of Ogg Vorbis."));
		}

		// Header type field indicates if this page is the beginning,
		// end, or middle of an Ogg logical bitstream.
		// It is permitted for a page to be both beginning and end - that means it's the only page.
		headerType = oggStream.ReadByte();
		UNUSED bool continuation = GetBit(headerType, 0);
		UNUSED bool beginningOfStream = GetBit(headerType, 1);
		bool endOfStream = GetBit(headerType, 2);

		// End of stream bit was set in the header type field - this is the last page of the logical bitstream
		if(endOfStream)
		{
			break;
		}
			
		// In Vorbis logical bitstreams, the granule position is the number of the last sample
		// contained in this frame.
		UNUSED int64_t granulePosition = oggStream.ReadInt64();
		// Bitstream serial number might be of interest if we wanted to be able to handle Ogg files with
		// multiple logical bitstreams...but we don't care.
		int32_t bitstreamSerialNumber = oggStream.ReadInt32();
		if(sampleRate != 0 && bitstreamSerialNumber != savedBitstreamSerialNumber)
		{
			throw OggVorbisException(ustr("The file is not a simple Ogg Vorbis file."));
		}
		savedBitstreamSerialNumber = bitstreamSerialNumber;

		UNUSED int32_t pageSequenceNumber = oggStream.ReadInt32();
		UNUSED int32_t checksum = oggStream.ReadInt32();
		uint8_t numSegments = oggStream.ReadByte();

		if(sampleRate == 0)
		{
			// This is the first Ogg page, so it's the primary Vorbis header.
			vector<uint8_t> segmentSizes = oggStream.ReadFull(numSegments);
			ogg_uint32_t vorbisHeaderPacketSize = 0;
			for(int segIndex = 0; segIndex < numSegments; segIndex++)
			{
				vorbisHeaderPacketSize += segmentSizes[segIndex];
				if(segmentSizes[segIndex] < 255)
				{
					break; // a segment size of less than 255 indicates the end of a packet.
				}
			}

			if(vorbisHeaderPacketSize < 16)
			{
				throw OggVorbisException(vorbisCorruptMsg);
			}

			uint8_t packetType = oggStream.ReadByte();
			if(packetType != 1)
			{
				throw OggVorbisException(vorbisCorruptMsg);
			}
				
			vector<uint8_t> vorbisString = oggStream.ReadFull(6);
			if(vorbisString[0] != 'v' || vorbisString[1] != 'o' || vorbisString[2] != 'r' 
			|| vorbisString[3] != 'b' || vorbisString[4] != 'i' || vorbisString[5] != 's')
			{
				throw OggVorbisException(vorbisCorruptMsg);
			}

			uint32_t vorbisVersion = oggStream.ReadUint32();
			if(vorbisVersion != 0)
			{
				throw OggVorbisException(vorbisCorruptMsg);
			}

			UNUSED uint8_t numChannels = oggStream.ReadByte();
			sampleRate = oggStream.ReadUint32();
			if(sampleRate == 0)
			{
				throw OggVorbisException(vorbisCorruptMsg);
			}

			int32_t pageDataSize = 0;
			for(uint8_t segmentIndex = 0; segmentIndex < numSegments; segmentIndex++)
			{
				uint8_t segmentSize = segmentSizes[segmentIndex];
				pageDataSize += segmentSize;
			}

			// Skip the rest of the page, we're not interested in it.
			int32_t unreadDataBytes = pageDataSize - 16;
			oggStream.Seek(unreadDataBytes, SeekOrigin::Current);
		}
		else
		{
			// We're not reading the data portion of the page, so we don't need to know how
			// large each segment is, just the total size.
			int32_t pageDataSize = 0;
			for(uint8_t segmentIndex = 0; segmentIndex < numSegments; segmentIndex++)
			{
				uint8_t segmentSize = oggStream.ReadByte();
				pageDataSize += segmentSize;
			}

			// Skip the data of the page, we're not interested in it.
			oggStream.Seek(pageDataSize, SeekOrigin::Current);
		}
			
	}

	if(sampleRate == 0)
	{
		throw OggVorbisException(vorbisCorruptMsg);
	}

	// Converting from seconds to samples might cause the result to be off be 1 if the number of seconds
	// came from GetRealTime().
	int64_t numSamples = static_cast<int64_t>(numSeconds * sampleRate);

	// Remember where the granule position field is, we're going to edit it later.
	// For now, we're reading the entire page so we can calculate what the checksum
	// should be after we change the granule position field.
	int64_t granulePositionPosition = oggStream.GetPosition();

	int64_t granulePosition = oggStream.ReadInt64();
	granulePosition = numSamples; // Set to what it will be for checksum calculation

	int32_t bitstreamSerialNumber = oggStream.ReadInt32();
	int32_t pageSequenceNumber = oggStream.ReadInt32();

	int32_t checksum = oggStream.ReadInt32();
	checksum = 0; // Set to 0 for checksum calculation

	uint8_t numSegments = oggStream.ReadByte();

	vector<uint8_t> segmentSizes = oggStream.ReadFull(numSegments);

	// Reconstruct the header (with the new granule position) as a byte vector for checksumming
	vector<uint8_t> headerBytes;
	headerBytes.push_back('O');
	headerBytes.push_back('g');
	headerBytes.push_back('g');
	headerBytes.push_back('S');
	headerBytes.push_back(version);
	headerBytes.push_back(headerType);
	AppendBytes(headerBytes, granulePosition);
	AppendBytes(headerBytes, bitstreamSerialNumber);
	AppendBytes(headerBytes, pageSequenceNumber);
	AppendBytes(headerBytes, checksum);
	headerBytes.push_back(numSegments);
	headerBytes.insert(headerBytes.end(), segmentSizes.begin(), segmentSizes.end());

		
	// Calculate the size of the data part of the page
	int32_t pageDataSize = 0;
	for(uint8_t segmentIndex = 0; segmentIndex < numSegments; segmentIndex++)
	{
		uint8_t segmentSize = segmentSizes[segmentIndex];
		pageDataSize += segmentSize;
	}

	// Read the entire data part of the page into a byte vector for checksumming
	vector<uint8_t> dataBytes = oggStream.ReadFull(pageDataSize);

	// Create an ogg_page structure using the header and body byte vectors so that libogg can do the checksum
	ogg_page page;
	page.header_len = headerBytes.size();
	page.header = &(headerBytes[0]);
	page.body_len = dataBytes.size();
	if(dataBytes.size() > 0)
	{
		page.body = &(dataBytes[0]);
	}

	// Let libogg do the tricky CRC stuff
	ogg_page_checksum_set(&page);
	checksum = GetFromBytes<int32_t>(&(headerBytes[0]), Endianness::Little, 22);

	// Finally, write the updated granule position and checksum. We're not changing the file
	// size or moving anything around, so we can just edit the file in place.
	oggStream.Seek(granulePositionPosition, SeekOrigin::Begin);
	oggStream.WriteInt64(granulePosition);
	oggStream.Seek(8, SeekOrigin::Current);
	oggStream.WriteInt32(checksum);
}

}

/*
 Copyright 2011 Greg Najda

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/
